#
# Ronin Exploits - A Ruby library for Ronin that provides exploitation and
# payload crafting functionality.
#
# Copyright (c) 2007-2013 Hal Brodigan (postmodern.mod3 at gmail.com)
#
# This file is part of Ronin Exploits.
#
# Ronin is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Ronin is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Ronin.  If not, see <http://www.gnu.org/licenses/>
#

require 'ronin/config'

require 'set'
require 'tempfile'

module Ronin
  module Exploits
    module Helpers
      #
      # Adds methods for building and managing files to be used with an
      # exploit.
      #
      # ## Parameters
      #
      # * `output_dir` (`Ronin::Config::TMP_DIR`) - Directory to save built
      #   file in.
      # * `output_file_name` (`'exploit'`) - Name of the file to build.
      # * `clean_output_files` (`true`) - Specifies whether the built files
      #   will be deleted upon exit.
      #
      module FileBuilder
        def self.extended(obj)
          obj.instance_eval do
            # The output directory for file creation
            parameter :output_dir,
                      type: String,
                      default: Ronin::Config::TMP_DIR,
                      description: 'Directory to save built file in'

            # The name of the output file
            parameter :output_file_name,
                      type: String,
                      default: 'exploit',
                      description: 'Name of the file to build'

            # Whether or not to delete the output file at exit
            parameter :clean_output_files,
                      type: true,
                      default: true,
                      description: 'Delete the file on exit'
          end
        end

        #
        # The paths of previously built files.
        #
        # @return [Set]
        #   The unique set of paths.
        #
        # @since 1.0.0
        #
        def output_file_paths
          @output_file_paths ||= Set[]
        end

        protected

        #
        # Opens the file to be built.
        #
        # @param [String] name
        #   Optional name of the file to build.
        #
        # @yield [file]
        #   If a block is given, it will be passed the newly opened file.
        #   After the block has returned, the file will be closed.
        #
        # @yieldparam [Tempfile, File] file
        #   The newly opened file.
        #
        # @return [String]
        #   Path to the built file.
        #
        # @example
        #   build_file do |file|
        #     file << 'some data'
        #   end
        #
        # @example Generate a specifically named file.
        #   build_file('file1.xml') do |file|
        #     file << "<evil />"
        #   end
        #
        def build_file(name=self.output_file_name,&block)
          path = nil

          if self.clean_output_files
            Tempfile.open(name,self.output_dir) do |file|
              path = file.path

              print_debug "Building file #{path} ..."
              block.call(file)
            end
          else
            name = File.expand_path(File.join('',name))
            path = File.join(self.output_dir,name)

            print_debug "Building file #{path} ..."
            File.open(self.output_file_path,'w',&block)
          end

          print_info "Built file #{path}"

          self.output_file_paths << path
          return path
        end
      end
    end
  end
end
